package cn.monkey.io.transport.mq.rabbitmq;

import cn.monkey.io.transport.core.FluxReceive;
import com.rabbitmq.client.AMQP;
import com.rabbitmq.client.Channel;
import com.rabbitmq.client.DefaultConsumer;
import com.rabbitmq.client.Envelope;

import java.io.IOException;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Stream;

public final class RabbitMqFluxReceive extends FluxReceive {

    final Collection<Channel> channels;
    final List<CompletableFuture<DefaultConsumer>> completableFutures;

    public RabbitMqFluxReceive(Collection<Channel> channels,
                               String... topics) {
        this.channels = channels;
        this.completableFutures = Stream.of(topics)
                .flatMap(topic -> Stream.of(this.channels)
                        .flatMap(Collection::stream)
                        .map(ch -> CompletableFuture.supplyAsync(() -> {
                            DefaultConsumer defaultConsumer = new DefaultConsumer(ch) {
                                @Override
                                public void handleDelivery(String consumerTag, Envelope envelope, AMQP.BasicProperties properties, byte[] body) throws IOException {
                                    RabbitMqFluxReceive.this.onInboundNext(body);
                                }
                            };
                            try {
                                ch.basicConsume(topic, defaultConsumer);
                            } catch (IOException e) {
                                super.error = e;
                                super.drainReceiver();
                            }
                            return defaultConsumer;
                        }))).toList();
    }

    @Override
    public void dispose() {
        this.completableFutures.forEach(f -> {
            f.whenCompleteAsync((consumer, e) -> {
                if (e != null) {
                    super.error = e;
                    super.drainReceiver();
                    return;
                }
                String consumerTag = consumer.getConsumerTag();
                if (consumerTag == null) {
                    return;
                }
                try {
                    consumer.getChannel().basicCancel(consumerTag);
                } catch (IOException ee) {
                    super.error = ee;
                    super.drainReceiver();
                }
            });
        });
        super.dispose();
    }
}
